"""DHF模型主模块

大伙房流域降雨径流预报模型的主要实现类，整合产流计算和汇流计算。
"""
import os
import json
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Any

import numpy as np
import pandas as pd
from loguru import logger
from matplotlib import pyplot as plt

from .dhf_utils import ModelUtils
from .parameters import DHFParameters
from .routing import Routing
from .runoff_generation import RunoffGeneration


class DHFModel:
    def __init__(self, initial_flow: float = 50.0,
                 rainfall_center: str = '中游',
                 watershed_area: float = 1000.0,
                 config_file: str = None):
        """
        初始化DHF模型

        Args:
            initial_flow: 降雨开始时的起始流量 (m³/s)
            rainfall_center: 降雨中心位置 ('上游', '中游', '下游')
            watershed_area: 流域面积 (km²)
            config_file: 参数配置文件路径（JSON格式）
        """
        # 从配置文件加载参数（如果提供）
        if config_file and Path(config_file).exists():
            logger.info(f"从配置文件加载参数: {config_file}")
            self.parameters = DHFParameters.from_json(config_file, initial_flow)
        else:
            # 使用默认参数
            self.parameters = DHFParameters(initial_flow)

        self.watershed_area = watershed_area
        self.rainfall_center = rainfall_center

        # 初始化计算模块
        self.runoff_generation = RunoffGeneration(self.parameters)
        self.routing = Routing(self.parameters, rainfall_center)
        self.routing.set_watershed_area(watershed_area)

        # 存储计算结果
        self.results = {}
        self.input_data = None

        # 模型状态
        self.is_calibrated = False
        self.simulation_completed = False

    def set_initial_conditions(self, surface_storage: float = None,
                               subsurface_storage: float = None,
                               groundwater_storage: float = None):
        """
        设置模型初始条件

        Args:
            surface_storage: 表层蓄水量初值 (mm)
            subsurface_storage: 下层蓄水量初值 (mm)
            groundwater_storage: 地下水库蓄水量初值 (mm)
        """
        self.runoff_generation.set_state_variables(
            S_a=surface_storage,
            U_a=subsurface_storage,
            V_a=groundwater_storage
        )

    def prepare_input_data(self, rainfall_data: List[float],
                           evaporation_data: List[float] = None,
                           time_step: float = 1.0,
                           start_time: datetime = None) -> pd.DataFrame:
        """
        准备模型输入数据

        Args:
            rainfall_data: 降雨数据序列 (mm)
            evaporation_data: 蒸发数据序列 (mm)
            time_step: 时间步长 (小时)
            start_time: 开始时间

        Returns:
            输入数据DataFrame
        """
        if start_time is None:
            start_time = datetime(2023, 7, 1)  # 默认开始时间

        # 如果没有提供蒸发数据，使用默认值
        if evaporation_data is None:
            # 根据月份和降雨量估算蒸发量
            evaporation_data = []
            for rainfall in rainfall_data:
                month = start_time.month
                evap = self.parameters.get_evaporation_capacity(month, rainfall)
                evaporation_data.append(evap)

        # 创建时间索引
        time_index = pd.date_range(
            start=start_time,
            periods=len(rainfall_data),
            freq=f'{int(time_step)}H'
        )

        self.input_data = pd.DataFrame({
            'rainfall': rainfall_data,
            'evaporation': evaporation_data,
            'time_step': time_step
        }, index=time_index)

        return self.input_data

    def run_simulation(self, rainfall_data: List[float],
                       evaporation_data: List[float] = None,
                       time_step: float = 1.0,
                       start_time: datetime = None) -> Dict[str, Any]:
        """
        运行DHF模型模拟

        Args:
            rainfall_data: 降雨数据序列 (mm)
            evaporation_data: 蒸发数据序列 (mm)
            time_step: 时间步长 (小时)
            start_time: 开始时间

        Returns:
            模拟结果字典
        """
        # 准备输入数据
        input_df = self.prepare_input_data(
            rainfall_data, evaporation_data, time_step, start_time
        )

        # 初始化结果存储
        runoff_components = {
            'impervious_runoff': [],
            'interflow': [],
            'groundwater_runoff': [],
            'total_runoff': [],
            'surface_storage': [],
            'subsurface_storage': []
        }

        # 在文件顶部添加loguru导入
        from loguru import logger

        # 在run_simulation方法中:
        # 逐时段计算产流
        logger.info("开始产流计算...")
        for i, (timestamp, row) in enumerate(input_df.iterrows()):
            rainfall = row['rainfall']
            evaporation = row['evaporation']

            # 计算产流分量
            components = self.runoff_generation.calculate_runoff_components(
                rainfall, evaporation, time_step
            )

            # 存储结果
            for key in runoff_components:
                if key in components:
                    runoff_components[key].append(components[key])

            # 显示进度
            if (i + 1) % max(1, len(input_df) // 10) == 0:
                progress = (i + 1) / len(input_df) * 100
                logger.debug(f"产流计算进度: {progress:.1f}%")

        # 汇流计算
        logger.info("开始汇流计算...")
        routed_components = self.routing.calculate_routing_for_components(
            runoff_components, time_step
        )

        # 计算总流量
        total_flow = self.routing.calculate_total_flow(routed_components)

        # 计算洪峰信息
        peak_info = self.routing.get_peak_flow_info(total_flow, time_step)

        # 计算径流总量
        total_volume = self.routing.calculate_runoff_volume(total_flow, time_step)

        # 存储结果
        self.results = {
            'input_data': input_df,
            'runoff_components': runoff_components,
            'routed_components': routed_components,
            'flow_series': total_flow,
            'peak_info': peak_info,
            'total_volume': total_volume,
            'time_series': input_df.index,
            'simulation_params': {
                'watershed_area': self.watershed_area,
                'rainfall_center': self.rainfall_center,
                'time_step': time_step
            }
        }

        self.simulation_completed = True
        logger.success("模拟计算完成！")

        return self.results

    def evaluate_performance(self, observed_flow: List[float]) -> Dict[str, float]:
        """
        评估模型性能

        Args:
            observed_flow: 观测流量序列 (m³/s)

        Returns:
            性能指标字典
        """
        if not self.simulation_completed:
            raise ValueError("请先运行模拟计算")

        simulated_flow = self.results['flow_series']
        performance = ModelUtils.calculate_model_performance(
            np.array(observed_flow), simulated_flow
        )

        self.results['performance'] = performance
        self.results['observed_flow'] = observed_flow

        return performance

    def calibrate_parameters(self, rainfall_data: List[float],
                             observed_flow: List[float],
                             evaporation_data: List[float] = None,
                             max_iterations: int = 100) -> Dict[str, Any]:
        """
        参数率定 (简化版本)

        Args:
            rainfall_data: 降雨数据序列 (mm)
            observed_flow: 观测流量序列 (m³/s)
            evaporation_data: 蒸发数据序列 (mm)
            max_iterations: 最大迭代次数

        Returns:
            率定结果字典
        """
        logger.info("开始参数率定...")

        best_nse = -np.inf
        best_params = None

        # 简化的参数率定：调整关键参数
        param_ranges = {
            'K': (0.1, 1.0),  # 下层下渗曲线曲率
            'B': (0.5, 2.0),  # 下渗抛物线分配曲线形状参数
            'S0': (20, 100),  # 表层蓄水容量
            'Ka': (0.5, 0.95)  # 前期净雨影响程度衰减系数
        }

        # 随机搜索法进行参数优化
        for iteration in range(max_iterations):
            # 随机生成参数组合
            test_params = {}
            for param, (min_val, max_val) in param_ranges.items():
                test_params[param] = np.random.uniform(min_val, max_val)

            # 更新参数
            original_params = {}
            for param, value in test_params.items():
                if param in self.parameters.runoff_params:
                    original_params[param] = self.parameters.runoff_params[param]
                    self.parameters.runoff_params[param] = value
                elif param in self.parameters.routing_params:
                    original_params[param] = self.parameters.routing_params[param]
                    self.parameters.routing_params[param] = value

            # 重新初始化计算模块
            self.runoff_generation = RunoffGeneration(self.parameters)
            self.routing = Routing(self.parameters, self.rainfall_center)
            self.routing.set_watershed_area(self.watershed_area)

            try:
                # 运行模拟
                results = self.run_simulation(rainfall_data, evaporation_data)

                # 计算性能指标
                performance = self.evaluate_performance(observed_flow)
                nse = performance['NSE']

                # 更新最佳参数
                if nse > best_nse:
                    best_nse = nse
                    best_params = test_params.copy()
                    logger.debug(f"迭代 {iteration + 1}: NSE = {nse:.4f} (最佳)")
                else:
                    # 恢复原参数
                    for param, value in original_params.items():
                        if param in self.parameters.runoff_params:
                            self.parameters.runoff_params[param] = value
                        elif param in self.parameters.routing_params:
                            self.parameters.routing_params[param] = value

            except Exception as e:
                # 恢复原参数
                for param, value in original_params.items():
                    if param in self.parameters.runoff_params:
                        self.parameters.runoff_params[param] = value
                    elif param in self.parameters.routing_params:
                        self.parameters.routing_params[param] = value
                    continue

            # 显示进度
            if (iteration + 1) % max(1, max_iterations // 10) == 0:
                progress = (iteration + 1) / max_iterations * 100
                logger.info(f"率定进度: {progress:.1f}%, 当前最佳NSE: {best_nse:.4f}")

        # 应用最佳参数
        if best_params:
            logger.info("应用最佳参数到模型...")
            for param, value in best_params.items():
                if param in self.parameters.runoff_params:
                    self.parameters.runoff_params[param] = value
                elif param in self.parameters.routing_params:
                    self.parameters.routing_params[param] = value

            # 重新初始化计算模块以使用新参数
            self.runoff_generation = RunoffGeneration(self.parameters)
            self.routing = Routing(self.parameters, self.rainfall_center)
            self.routing.set_watershed_area(self.watershed_area)

            # 使用最佳参数重新运行模拟
            final_results = self.run_simulation(rainfall_data, evaporation_data)
            final_performance = self.evaluate_performance(observed_flow)

            # 自动保存最佳参数
            os.makedirs("./output", exist_ok=True)
            self.save_current_parameters("./output/calibrated_parameters.json")

            logger.success("最佳参数已应用到模型并保存")

            self.is_calibrated = True

            calibration_results = {
                'best_parameters': best_params,
                'best_nse': best_nse,
                'final_performance': final_performance,
                'iterations': max_iterations
            }

            logger.success(f"参数率定完成！最佳NSE: {best_nse:.4f}")
            return calibration_results

        else:
            logger.warning("参数率定失败，未找到改进的参数组合")
            return {'best_parameters': None, 'best_nse': best_nse}

    def generate_plots(self, save_dir: str = None) -> Dict[str, plt.Figure]:
        """
        生成模型结果图表

        Args:
            save_dir: 图表保存目录

        Returns:
            图表对象字典
        """
        if not self.simulation_completed:
            raise ValueError("请先运行模拟计算")

        plots = {}
        time_series = self.results['time_series']
        flow_series = self.results['flow_series']

        # 确保时间序列和流量序列长度匹配
        min_length = min(len(time_series), len(flow_series))
        time_series_aligned = time_series[:min_length]
        flow_series_aligned = flow_series[:min_length]

        # 流量过程线
        observed = self.results.get('observed_flow')
        if observed is not None:
            observed = observed[:min_length]  # 同样截取观测数据

        rainfall = self.results['input_data']['rainfall'].values
        rainfall_aligned = rainfall[:min_length]  # 截取降雨数据

        plots['hydrograph'] = ModelUtils.plot_hydrograph(
            time_series_aligned, observed, flow_series_aligned, rainfall_aligned,
            title="DHF模型流量过程线",
            save_path=f"{save_dir}/hydrograph.png" if save_dir else None
        )

        # 径流分量过程线
        rainfall_data = self.results['input_data']['rainfall'].values
        rainfall_aligned = rainfall_data[:min_length]  # 截取降雨数据

        plots['components'] = ModelUtils.plot_runoff_components(
            time_series_aligned, self.results['runoff_components'],
            rainfall_data=rainfall_aligned,  # 添加降雨数据参数
            title="DHF模型径流分量过程线",
            save_path=f"{save_dir}/components.png" if save_dir else None
        )

        # 蓄水量动态变化
        plots['storage'] = ModelUtils.plot_storage_dynamics(
            time_series,
            np.array(self.results['runoff_components']['surface_storage']),
            np.array(self.results['runoff_components']['subsurface_storage']),
            title="DHF模型蓄水量动态变化",
            save_path=f"{save_dir}/storage.png" if save_dir else None
        )

        # 观测值与模拟值对比散点图
        if observed is not None:
            plots['scatter'] = ModelUtils.plot_scatter_comparison(
                np.array(observed), flow_series_aligned,  # ✅ 使用正确的变量
                title="DHF模型观测值与模拟值对比",
                save_path=f"{save_dir}/scatter.png" if save_dir else None
            )

        return plots

    def export_results(self, file_path: str):
        """
        导出模型结果到Excel文件

        Args:
            file_path: 文件保存路径
        """
        if not self.simulation_completed:
            raise ValueError("请先运行模拟计算")

        ModelUtils.export_results_to_excel(
            self.results, file_path, self.results['time_series']
        )

    def generate_report(self) -> str:
        """
        生成模型运行报告

        Returns:
            报告字符串
        """
        if not self.simulation_completed:
            raise ValueError("请先运行模拟计算")

        return ModelUtils.create_model_summary_report(self.results)

    def get_model_info(self) -> Dict[str, Any]:
        """
        获取模型基本信息

        Returns:
            模型信息字典
        """
        return {
            'model_name': 'DHF (大伙房) 水文模型',
            'version': '1.0.0',
            'watershed_area': self.watershed_area,
            'rainfall_center': self.rainfall_center,
            'initial_flow': self.parameters.initial_flow,
            'is_calibrated': self.is_calibrated,
            'simulation_completed': self.simulation_completed,
            'parameters': self.parameters.get_all_parameters()
        }

    def reset_model(self):
        """
        重置模型状态
        """
        # 重新初始化计算模块
        self.runoff_generation = RunoffGeneration(self.parameters)
        self.routing = Routing(self.parameters, self.rainfall_center)
        self.routing.set_watershed_area(self.watershed_area)

        # 清空结果
        self.results = {}
        self.input_data = None
        self.simulation_completed = False

        logger.info("模型状态已重置")

    def __str__(self) -> str:
        """模型信息的字符串表示"""
        return f"DHF水文模型 (流域面积: {self.watershed_area}km², 降雨中心: {self.rainfall_center})"

    def __repr__(self) -> str:
        return self.__str__()

    def load_parameters_from_json(self, json_path: str):
        """
        从JSON文件重新加载参数

        Args:
            json_path: JSON参数文件路径
        """
        logger.info(f"重新加载参数: {json_path}")
        self.parameters = DHFParameters.from_json(json_path, self.parameters.initial_flow)

        # 重新初始化计算模块
        self.runoff_generation = RunoffGeneration(self.parameters)
        self.routing = Routing(self.parameters, self.rainfall_center)
        self.routing.set_watershed_area(self.watershed_area)

        # 重置状态
        self.is_calibrated = False
        self.simulation_completed = False

        logger.success("参数重新加载完成")

    def save_current_parameters(self, json_path: str):
        """
        保存当前参数到JSON文件

        Args:
            json_path: JSON文件保存路径
        """
        self.parameters.save_to_json(json_path)
